package com.lfy.modules.ai.controller;

import cn.hutool.core.date.DateUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.lfy.modules.ai.cache.AiDispatchCache;
import com.lfy.modules.ai.domain.AiModelEntity;
import com.lfy.modules.ai.domain.AiModelUserContantEntity;
import com.lfy.modules.ai.domain.AiOwnEntity;
import com.lfy.modules.ai.domain.dto.CreateAiUserDTO;
import com.lfy.modules.ai.domain.dto.ModifyAiModelUseDTO;
import com.lfy.modules.ai.domain.dto.ModifyAiUserUseDTO;
import com.lfy.modules.ai.domain.vo.AiModelUserVO;
import com.lfy.modules.ai.domain.vo.AiUseVO;
import com.lfy.modules.ai.mapper.AiModelUserContantMapper;
import com.lfy.modules.ai.service.AiModelService;
import com.lfy.modules.ai.service.AiModelUserContantService;
import com.lfy.modules.ai.service.AiOwnService;
import com.lfy.modules.blog.domain.RoleEntity;
import com.lfy.modules.blog.service.RoleService;
import com.lfy.modules.blog.service.UserService;
import com.lfy.modules.common.anno.Authentication;
import com.lfy.modules.common.anno.AutoUser;
import com.lfy.modules.common.custom.constants.LockPrefixConstant;
import com.lfy.modules.common.domain.PageResultVO;
import com.lfy.modules.common.domain.R;
import com.lfy.modules.common.domain.entity.UserEntity;
import com.lfy.modules.common.utils.PageUtil;
import com.lfy.modules.common.utils.RequestUtil;
import com.lfy.modules.letter.domain.SiteLetterEntity;
import com.lfy.modules.letter.domain.builder.Letter;
import com.lfy.modules.letter.domain.builder.LetterEnum;
import com.lfy.modules.letter.domain.event.LetterEvent;
import org.apache.catalina.User;
import org.springframework.beans.BeanUtils;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;

import javax.annotation.Resource;
import java.util.*;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.function.Function;
import java.util.stream.Collectors;

@RestController
@RequestMapping("admin/ai/dispatch")
public class AiDispatchController {
    @Resource
    private AiOwnService aiOwnService;
    @Resource
    private AiModelService aiModelService;

    @Resource
    private AiModelUserContantService aiModelUserContantService;
    @Resource
    private AiModelUserContantMapper aiModelUserContantMapper;
    @Resource
    private UserService userService;
    @Resource
    private RoleService roleService;
    @Resource
    private AiDispatchCache aiDispatchCache;
    @Resource
    private ThreadPoolExecutor mainThreadPool;
    /**
     * 获取用户列表
     */
    @GetMapping("/sys/user/list")
    public R<List<UserEntity>> getSysUserList(@RequestParam(value = "modelId") Integer modelId) {
        List<AiModelUserContantEntity> list = this.aiModelUserContantMapper.selectList(Wrappers.lambdaQuery(AiModelUserContantEntity.class)
                .eq(AiModelUserContantEntity::getAiModelId, modelId));
        List<UserEntity> userList = userService.list(Wrappers.lambdaQuery(UserEntity.class).select(UserEntity::getId, UserEntity::getUsername,UserEntity::getName));
        if(CollectionUtils.isEmpty(list)){
            return R.success(userList);
        }
        Set<Long> ids = list.stream().map(AiModelUserContantEntity::getUserId).collect(Collectors.toSet());
        List<UserEntity> result = userList.stream().filter(item -> !ids.contains(item.getId())).collect(Collectors.toList());
        return R.success(result);
    }
    /**
     * 获取模型列表
     */
    @GetMapping("/get/list")
    @Authentication("ai:get:list")
    public R<List<AiOwnEntity>> getAiList() {
        List<AiOwnEntity> list = aiOwnService.list(Wrappers.lambdaQuery(AiOwnEntity.class));
        return R.success(list);
    }
    /**
     * 获取模型使用情况
     * @param id 模型id
     * @return 模型使用情况
     */
    @GetMapping("/get/use")
    @Authentication("ai:get:use")
    public R<List<AiUseVO>> getAiUse(@RequestParam("id") Integer id){
        List<AiModelEntity> list = aiModelService.list(Wrappers.lambdaQuery(AiModelEntity.class).eq(AiModelEntity::getAiId, id));
        if(CollectionUtils.isEmpty(list)){
            return R.error("该模型尚未开通");
        }

        List<AiUseVO> result = list.stream().map(item -> {

            AiUseVO aiUseVO = new AiUseVO();
            aiUseVO.setId(item.getId());
            aiUseVO.setTitle(item.getAiModelName());
            aiUseVO.setRemaining(item.getAiRemaining());
            aiUseVO.setAiLock(item.getAiLock());
            return aiUseVO;
        }).collect(Collectors.toList());
        return R.success(result);
    }
    /**
     * 扩缩容
     */

    @PostMapping("/update/model")
    @Authentication("ai:update:model")
    public R<String> updateModel(@RequestBody @Validated ModifyAiModelUseDTO dto, @AutoUser UserEntity userEntity) {
        synchronized ((LockPrefixConstant.AI_MODEL_LOCK_PREFIX + dto.getModelId()).intern()){
            AiModelEntity one = this.aiModelService.getById(dto.getModelId());
            if (Objects.isNull(one)) {
                return R.error("该模型不存在");
            }
            if(one.getAiRemaining() + dto.getNum() < 0){
                return R.error("资源不能小于0");
            }
            one.setAiRemaining(one.getAiRemaining() + dto.getNum());
            one.setAiGrossCapacity(one.getAiGrossCapacity() + dto.getNum());
            aiModelService.updateById(one);
            return R.success("扩缩容成功");
        }
    }

    /**
     * 开通用户资源
     */
    @PostMapping("/create/user/")
    @Authentication("ai:create:user")
    @Transactional(rollbackFor = Exception.class)
    public R<String> createUser(@RequestBody @Validated CreateAiUserDTO dto, @AutoUser UserEntity userEntity) {
        Integer aiModel = dto.getModelId();
        AiModelEntity model = this.aiModelService.getById(aiModel);
        if(Objects.isNull(model)){
            return R.error("该模型不存在");
        }
        if(model.getAiLock() >= model.getAiRemaining()){
            return R.error("分配资源已达上限");
        }
        AiModelUserContantEntity one = this.aiModelUserContantService.getOne(Wrappers.lambdaQuery(AiModelUserContantEntity.class)
                .eq(AiModelUserContantEntity::getAiModelId, aiModel).eq(AiModelUserContantEntity::getUserId, dto.getUserId()));
        if(Objects.nonNull(one)){
            return R.error("该用户已开通该模型");
        }
        AiModelUserContantEntity aiModelUserContantEntity = new AiModelUserContantEntity();
        aiModelUserContantEntity.setCreateUser(userEntity.getId());
        aiModelUserContantEntity.setAiModelId(dto.getModelId());
        aiModelUserContantEntity.setCapacityUsed(0L);
        aiModelUserContantEntity.setCapacityAll(dto.getCapacityAll());
        aiModelUserContantEntity.setAiId(model.getAiId());
        aiModelUserContantEntity.setUserId(dto.getUserId());
        aiModelUserContantEntity.setAiModelKey(model.getAiModelKey());

        String ip = RequestUtil.getIp();

        aiModelUserContantEntity.setCreateIp(ip);
        aiModelUserContantEntity.setCreateIpAddress(RequestUtil.getPlace(ip));
        aiModelUserContantService.save(aiModelUserContantEntity);
        // 更新模型使用量
        model.setAiLock(model.getAiLock() + dto.getCapacityAll());
        model.setAiRemaining(model.getAiRemaining() - dto.getCapacityAll());
        aiModelService.updateById(model);
        SpringUtil.publishEvent(new LetterEvent(this, Letter.builder(LetterEnum.LetterType.SYS_SEND.getCode(),
                "AI模型开通通知","以为您开通"+model.getAiModelName()+"模型资源，请及时查看使用情况。",dto.getUserId(),LetterEnum.SendEmailType.NO.getCode())));
        mainThreadPool.execute(() -> {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            aiDispatchCache.initAiDispatchCache();
        });
        return R.success("开通成功");
    }
    /**
     * 删除用户资源
     */
    @PostMapping("/delete/user/{id}")
    @Authentication("ai:delete:user")
    @Transactional(rollbackFor = Exception.class)
    public R<String> deleteUser(@PathVariable("id") Integer id , @AutoUser UserEntity userEntity) {
        AiModelUserContantEntity one = this.aiModelUserContantService.getById(id);
        if(Objects.isNull(one)){
            return R.error("该用户资源不存在");
        }
        AiModelEntity model = this.aiModelService.getById(one.getAiModelId());
        if(Objects.isNull(model)){
            return R.error("该模型不存在");
        }
        // 更新模型使用量
        model.setAiLock(model.getAiLock() - one.getCapacityAll());
        model.setAiRemaining(model.getAiRemaining() + one.getCapacityAll());
        aiModelService.updateById(model);
        aiModelUserContantService.removeById(id);
        mainThreadPool.execute(() -> {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            aiDispatchCache.initAiDispatchCache();
        });
        return R.success("删除成功");

    }
    /**
     * 获取用户资源列表
     */
    @GetMapping("/get/user/list")
    @Authentication("ai:get:user:list")
    public R<PageResultVO<AiModelUserVO>> getUserList(@RequestParam("pageSize") Integer pageSize,
                                                      @RequestParam("pageNum") Integer pageNum,
                                                      @RequestParam(value = "aiId", required = false) Integer aiId,
                                                      @AutoUser UserEntity userEntity) {
        LambdaQueryWrapper<AiModelUserContantEntity> wrapper = Wrappers.lambdaQuery(AiModelUserContantEntity.class);
        if(Objects.nonNull(aiId)){
            wrapper.eq(AiModelUserContantEntity::getAiId, aiId);
        }
        Page<AiModelUserContantEntity> aiModelUserContantEntityPage = this.aiModelUserContantMapper.selectPage(new Page<AiModelUserContantEntity>(pageNum, pageSize), wrapper);
        List<AiModelUserContantEntity> list = aiModelUserContantEntityPage.getRecords();
        if(CollectionUtils.isEmpty(list)){
            return R.success(PageUtil.convert2PageEmptyResult(aiModelUserContantEntityPage, AiModelUserVO.class));
        }
        List<UserEntity> userEntities = this.userService.list();
        if(CollectionUtils.isEmpty(userEntities)){
            return R.error("用户列表为空");
        }
        List<RoleEntity> roleEntities = this.roleService.list();
        Map<Long, UserEntity> userMap = userEntities.stream().collect(Collectors.toMap(UserEntity::getId, Function.identity()));
        List<AiOwnEntity> aiOwnEntityList = this.aiOwnService.list();
        List<AiModelEntity> modelList = this.aiModelService.list();
        List<AiModelUserVO> result = list.stream().map(item -> {
            AiModelUserVO aiModelUserVO = new AiModelUserVO();
            BeanUtils.copyProperties(item, aiModelUserVO);
            aiModelUserVO.setCreateTime(DateUtil.formatDate(item.getCreateTime()));
            aiModelUserVO.setUpdateTime(DateUtil.formatDate(item.getUpdateTime()));
            UserEntity user = userMap.get(item.getUserId());
            if(Objects.nonNull(user)){
                aiModelUserVO.setUserName(user.getUsername());
            }
            UserEntity createUser = userMap.get(item.getCreateUser());
            if(Objects.nonNull(createUser)){
                aiModelUserVO.setCreateUserName(createUser.getUsername());
                roleEntities.stream().filter(role->role.getId().equals(createUser.getRole())).findFirst().ifPresent(role->
                        aiModelUserVO.setCreateUserRole(role.getRoleName()));
            }
            if(!CollectionUtils.isEmpty(aiOwnEntityList)){
                aiOwnEntityList.stream().filter(aiOwn -> aiOwn.getId().equals(item.getAiId())).findFirst().ifPresent(ai->
                        aiModelUserVO.setAiName(ai.getAiName()));
            }
            if(!CollectionUtils.isEmpty(modelList)){
                modelList.stream().filter(model -> model.getId().equals(item.getAiModelId())).findFirst().ifPresent(modelEntity -> {
                    aiModelUserVO.setAiModelName(modelEntity.getAiModelName());
                });
            }
            return aiModelUserVO;
        }).collect(Collectors.toList());
        return R.success(PageUtil.buildPageResultVO(aiModelUserContantEntityPage,result));
    }
    /**
     * 用户资源调整
     */
    @PostMapping("/sys/update/user")
    @Authentication("ai:sys:update:user")
    @Transactional(rollbackFor = Exception.class)
    public R<String> sysUpdateUser(@RequestBody @Validated ModifyAiUserUseDTO dto, @AutoUser UserEntity userEntity) {
        if(dto.getCapacityAll() <0){
            return R.error("资源不能小于0");
        }
        AiModelUserContantEntity one = aiModelUserContantService.getOne(Wrappers.lambdaQuery(AiModelUserContantEntity.class)
                .eq(AiModelUserContantEntity::getAiModelId, dto.getModelId()).eq(AiModelUserContantEntity::getUserId, dto.getUserId()));

        if(Objects.isNull(one)){
            return R.error("该用户资源不存在");
        }
        AiModelEntity modelEntity = this.aiModelService.getById(dto.getModelId());

        if(Objects.isNull(modelEntity)){
            return R.error("该模型不存在");
        }
        if(dto.getCapacityAll() > modelEntity.getAiRemaining()){
            return R.error("资源超出资源池上限");
        }
        //调整后
        Long keYong = one.getCapacityAll();
        //调整前
        Long capacityAll = dto.getCapacityAll();
        //算差值
        Long cha = capacityAll - keYong;
        //更新模型使用量
        modelEntity.setAiLock(modelEntity.getAiLock() + cha);
        modelEntity.setAiRemaining(modelEntity.getAiRemaining() - cha);
        this.aiModelService.updateById(modelEntity);
        //更新用户资源
        one.setCapacityAll(dto.getCapacityAll());
        this.aiModelUserContantService.updateById(one);
        return R.success("调整成功");
    }
}
